package qmlt.learning.decisiontree.control;

import java.util.Collection;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;

import qmlt.dataset.Instance;
import qmlt.dataset.utils.DataSetUtils;
import qmlt.dataset.utils.EntropyUtils;
import qmlt.dataset.utils.ListUtils;
import qmlt.learning.decisiontree.Node;

public class InfoGainTreeBuildController implements TreeBuildController
{
    @Override
    public void split(Node node)
    {
        if (node.instances.size() <= 0)
            return;

        Instance one = node.instances.get(0);
        Map<Object, List<Instance>> map0 = DataSetUtils.asMap(node.instances);
        if (map0.size() == 1)
        {
            node.targetClass = one.getTarget();
            return;
        }
        int dim = one.getFeatures().size();

        // find the best attribute
        int bestAttributeIndex = 0;
        float bestInfoGain = 0;
        Map<Object, List<Instance>> bestSplitting = null;
        for (int i = 0; i < dim; ++i)
        {
            Map<Object, List<Instance>> map = DataSetUtils.asMap(node.instances, i);

            List<Object> setOfClass = DataSetUtils.getFeatureColumn(node.instances, i);
            float entropy = EntropyUtils.calculateEntropy(setOfClass);
            float infoGain = getSplitScore(node.instances.size(), entropy, map.values());
            if (infoGain > bestInfoGain)
            {
                bestInfoGain = infoGain;
                bestAttributeIndex = i;
                bestSplitting = map;
            }
        }

        if (bestSplitting == null)
        {
            // can't get any info gain, make it a leaf and use the major class
            node.targetClass = DataSetUtils.findMajorClass(node.instances);
        }
        else
        {
            node.splittingAttributeIndex = bestAttributeIndex;
            node.children = new HashMap<Object, Node>();
            for (Object value : bestSplitting.keySet())
            {
                Node child = new Node();
                child.instances = bestSplitting.get(value);
                node.children.put(value, child);
            }
        }
    }

    @SuppressWarnings("unchecked")
    protected float getSplitScore(int size, float entropy, Collection<List<Instance>> subsets)
    {
        int n = subsets.size();
        List[] lists = new List[n];
        Iterator<List<Instance>> it = subsets.iterator();
        for (int i = 0; i < n; ++i)
        {
            lists[i] = DataSetUtils.getTargetColumn(it.next());
        }
        return EntropyUtils.calculateInfoGain(size, entropy, lists);
    }
}
